Changes

The approach here is quite straightfoward, _quantize_weights works as usual for 2D weights. The difference is in calculate hessian where the hessian is 3D in both 3D and 2D weights case. By default hessian has the shape (1, hidden_dim, hidden_dim).
Before this was just (hidden_dim, hidden_dim). For 3D, it is (num_experts/batch, hidden_dim, hidden_dim).

Now, this 3D hessian or "batched" hessian is looped over and the 2D weight is extracted and passed to the old _quantize_weights function as usual and scale/zp are returned. These scales and zp are then stacked together in a collector variable. For 2D case, it is flattened. For 3D the stacked scale, zp are returned.

NOTE: Scale Estimation + GPTQ support is not added for 3D weights yet

Reason for changes

Support 3D weights for models like MoE in GPTQ

Related tickets

175789 & 175212

Tests